{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook is intended to demonstrate how vessel segmentation methods of ITKTubeTK can be applied to multi-channel MRI (MRA + T1, T2, etc)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import itk\n",
    "from itk import TubeTK as ttk\n",
    "\n",
    "from itkwidgets import view\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "ImageType = itk.Image[itk.F, 3]\n",
    "ReaderType = itk.ImageFileReader[ImageType]\n",
    "ResampleType = ttk.ResampleImage[ImageType]\n",
    "\n",
    "#IMPORTANT: This assumed the data has been brain-stripped and resampled into isotropic spacing!!!\n",
    "reader = ReaderType.New(FileName=\"MRA-Iso.mha\")\n",
    "reader.Update()\n",
    "im1iso = reader.GetOutput()\n",
    "reader = ReaderType.New(FileName=\"MRA-Brain.mha\")\n",
    "reader.Update()\n",
    "im1Brain = reader.GetOutput()\n",
    "\n",
    "reader = ReaderType.New(FileName=\"MRT1-Iso.mha\")\n",
    "reader.Update()\n",
    "im2iso = reader.GetOutput()\n",
    "reader = ReaderType.New(FileName=\"MRT1-Brain.mha\")\n",
    "reader.Update()\n",
    "im2Brain = reader.GetOutput()\n",
    "\n",
    "reader = ReaderType.New(FileName=\"MRT2-Iso.mha\")\n",
    "reader.Update()\n",
    "im3iso = reader.GetOutput()\n",
    "reader = ReaderType.New(FileName=\"MRT2-Brain.mha\")\n",
    "reader.Update()\n",
    "im3Brain = reader.GetOutput()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "899b16c9c93b46b0b509a675495431bd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Viewer(geometries=[], gradient_opacity=0.22, point_sets=[], rendered_image=<itk.itkImagePython.itkImageF3; pro…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "view(im1iso)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-15.9078321  -22.9855871    6.66245282]\n",
      " [ -2.60888692  -2.65183671  -1.04874474]\n",
      " [-13.41554095 -12.13042089  -0.80657562]\n",
      " [-27.82812546 -24.4508738   18.37093524]\n",
      " [-42.88031365 -15.51437618  22.45765273]\n",
      " [-12.70866419  -3.12293465  14.04063522]\n",
      " [ 14.92985178  10.67370493  17.04517199]\n",
      " [-18.34659598   8.88268063  16.25303385]\n",
      " [-23.14955895  -7.12923827  -6.79776154]\n",
      " [-53.24560825 -10.77376713  26.09933615]]\n"
     ]
    }
   ],
   "source": [
    "imMath = ttk.ImageMath[ImageType,ImageType].New()\n",
    "imMath.SetInput(im1Brain)\n",
    "imMath.Blur(1.0)\n",
    "imBlur = imMath.GetOutput()\n",
    "imBlurArray = itk.GetArrayViewFromImage(imBlur)\n",
    "\n",
    "numSeeds = 10\n",
    "seedCoverage = 20\n",
    "seedCoord = np.zeros([numSeeds,3])\n",
    "for i in range(numSeeds):\n",
    "    seedCoord[i] = np.unravel_index(np.argmax(imBlurArray, axis=None), imBlurArray.shape)\n",
    "    indx = [int(seedCoord[i][0]),int(seedCoord[i][1]),int(seedCoord[i][2])]\n",
    "    minX = max(indx[0]-seedCoverage,0)\n",
    "    maxX = max(indx[0]+seedCoverage,imBlurArray.shape[0])\n",
    "    minY = max(indx[1]-seedCoverage,0)\n",
    "    maxY = max(indx[1]+seedCoverage,imBlurArray.shape[1])\n",
    "    minZ = max(indx[2]-seedCoverage,0)\n",
    "    maxZ = max(indx[2]+seedCoverage,imBlurArray.shape[2])\n",
    "    imBlurArray[minX:maxX,minY:maxY,minZ:maxZ]=0\n",
    "    indx.reverse()\n",
    "    seedCoord[:][i] = im1Brain.TransformIndexToPhysicalPoint(indx)\n",
    "print(seedCoord)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "**** Processing seed 0 : [-15.9078321  -22.9855871    6.66245282]\n",
      "**** Processing seed 1 : [-2.60888692 -2.65183671 -1.04874474]\n",
      "**** Processing seed 2 : [-13.41554095 -12.13042089  -0.80657562]\n",
      "**** Processing seed 3 : [-27.82812546 -24.4508738   18.37093524]\n",
      "**** Processing seed 4 : [-42.88031365 -15.51437618  22.45765273]\n",
      "**** Processing seed 5 : [-12.70866419  -3.12293465  14.04063522]\n",
      "**** Processing seed 6 : [14.92985178 10.67370493 17.04517199]\n",
      "**** Processing seed 7 : [-18.34659598   8.88268063  16.25303385]\n",
      "**** Processing seed 8 : [-23.14955895  -7.12923827  -6.79776154]\n",
      "**** Processing seed 9 : [-53.24560825 -10.77376713  26.09933615]\n"
     ]
    }
   ],
   "source": [
    "# Manually extract a few vessels to form an image-specific training set\n",
    "vSeg = ttk.SegmentTubes[ImageType].New()\n",
    "vSeg.SetInput(im1iso)\n",
    "vSeg.SetVerbose(True)\n",
    "vSeg.SetMinRoundness(0.1)\n",
    "vSeg.SetMinCurvature(0.001)\n",
    "vSeg.SetRadiusInObjectSpace( 1 )\n",
    "for i in range(numSeeds):\n",
    "    print(\"**** Processing seed \" + str(i) + \" : \" + str(seedCoord[i]))\n",
    "    vSeg.ExtractTubeInObjectSpace( seedCoord[i], i )\n",
    "    \n",
    "tubeMaskImage = vSeg.GetTubeMaskImage()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0c0c0a82b2c64c87bf5a54bf2a4ae857",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Viewer(geometries=[], gradient_opacity=0.22, point_sets=[], rendered_image=<itk.itkImagePython.itkImageF3; pro…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "view(tubeMaskImage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "LabelMapType = itk.Image[itk.UC,3]\n",
    "\n",
    "trMask = ttk.ComputeTrainingMask[ImageType,LabelMapType].New()\n",
    "trMask.SetInput( tubeMaskImage )\n",
    "trMask.SetGap( 3 )\n",
    "trMask.SetObjectWidth( 1 )\n",
    "trMask.SetNotObjectWidth( 1 )\n",
    "trMask.Update()\n",
    "fgMask = trMask.GetOutput()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "359438cf7a5443c7b3308285f70b7e3c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Viewer(geometries=[], gradient_opacity=0.22, point_sets=[], rendered_image=<itk.itkImagePython.itkImageUC3; pr…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "view(fgMask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "enhancer = ttk.EnhanceTubesUsingDiscriminantAnalysis[ImageType,LabelMapType].New()\n",
    "enhancer.SetInput( im1iso )\n",
    "enhancer.AddInput( im2iso )\n",
    "enhancer.AddInput( im3iso )\n",
    "enhancer.SetLabelMap( fgMask )\n",
    "enhancer.SetRidgeId( 255 )\n",
    "enhancer.SetBackgroundId( 128 )\n",
    "enhancer.SetUnknownId( 0 )\n",
    "enhancer.SetTrainClassifier(True)\n",
    "enhancer.SetUseIntensityOnly(True)\n",
    "enhancer.SetScales([0.3333,1,2.5])\n",
    "enhancer.Update()\n",
    "enhancer.ClassifyImages()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2dd468f87f654c5fa703f2c74627aea3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Viewer(geometries=[], gradient_opacity=0.22, point_sets=[], rendered_image=<itk.itkImagePython.itkImageF3; pro…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#view(enhancer.GetOutput())\n",
    "view(enhancer.GetClassProbabilityImage(0))\n",
    "#imMath = ttk.ImageMath[ImageType,ImageType].New(Input = segmenter.GetClassProbabilityImage(0))\n",
    "#imMath.AddImages( segmenter.GetClassProbabilityImage(1), 1, -1 )\n",
    "#view(imMath.GetOutput())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "#itk.imwrite( enhancer.GetClassProbabilityImage(0), \"MRA-VesselProb.mha\", compression=True)\n",
    "#itk.imwrite( enhancer.GetClassProbabilityImage(1), \"MRA-NotVesselProb.mha\", compression=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "im1vess = itk.SubtractImageFilter( Input1=enhancer.GetClassProbabilityImage(0), Input2=enhancer.GetClassProbabilityImage(1))\n",
    "itk.imwrite( im1vess, \"MRA-VesselEnhanced.mha\", compression=True)\n",
    "\n",
    "brainMask = itk.imread( \"MRMask-Brain.mha\", itk.F )\n",
    "im1BrainVess = itk.MultiplyImageFilter(Input1 = im1vess, Input2=brainMask)\n",
    "itk.imwrite( im1BrainVess, \"MRA-Brain-VesselEnhanced.mha\", compression=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
